#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2023 reinovo, Inc. All Rights Reserved 
#
# @Time    : 2023/12/28 下午5:13
# @Author  : hmm
# @Email   : liuyuhang0531@foxmail.com
# @File    : Apriltag_Detection.py
# @Software: PyCharm

import cv2 as cv
import numpy as np
import apriltag

class apriltag_detection(object):
    def __init__(self,tag_size, intrinsic_matrix, families='tag36h10'):
        """

        :param tag_size: 二维码尺寸
        :param intrinsic_matrix: 相机内参矩阵
        :param families: 检测的二维码字典
        """
        self.detector = apriltag.Detector(apriltag.DetectorOptions(families=families))
        self.mtx = intrinsic_matrix
        self.tag_size = tag_size

    def detection(self,image):
        """
        检测二维码
        :param image: 检测图像
        :return: 输出检测结果
        """
        gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)
        tags = self.detector.detect(gray)
        return tags

    def get_point(self,tags):
        """
        获取真实世界二维码位置
        :param tags: 检测结果
        :return: 输出坐标和id
        """
        data = []
        mtx = self.mtx
        cam_params = [mtx[0,0],mtx[1,1],mtx[0,2],mtx[1,2]]
        for tag in tags:
            M, e1, e2 = self.detector.detection_pose(tag, cam_params, self.tag_size)
            cam_point = tag.center
            img_point = np.linalg.inv(mtx) @ np.insert(cam_point, 2, 1)
            obj_point = M @ np.insert(img_point, 2, 1)
            data.append((tag.tag_id,obj_point[:3]))
        return dict(data)

    def darw(self,image, tags):
        """
        渲染检测结果
        :param image: 渲染图像
        :param tags: 检测结果
        :return: 渲染后图像
        """
        pos = np.array([[self.tag_size, 0, 0, 1],
                        [0, self.tag_size, 0, 1],
                        [0, 0, self.tag_size, 1]])

        color = [(0, 0, 255), (0, 255, 0), (255, 0, 0)]
        mtx = self.mtx
        cam_params = [mtx[0,0],mtx[1,1],mtx[0,2],mtx[1,2]]
        for tag in tags:
            M, e1, e2 = self.detector.detection_pose(tag, cam_params, self.tag_size)
            for i in range(len(pos)):
                axis = np.column_stack((self.mtx, [0, 0, 0])) @ M @ pos[i].T
                axis = axis / axis[2]
                cv.line(image, tuple(tag.center.astype(int)), tuple(axis[:2].astype(int)), color[i], 4)
            font = cv.FONT_HERSHEY_SIMPLEX
            cv.putText(image, f'tag_id:{tag.tag_id}', tuple(tag.center.astype(int)), font, 1, (248, 86, 105), 4, cv.LINE_AA)
        return image


if __name__ == "__main__":
    mtx = np.load('./data/camera_param.npz')['mtx']
    at_detect = apriltag_detection(30,mtx)
    cap = cv.VideoCapture(4)
    while cap.isOpened():
        ret,frame = cap.read()
        tags = at_detect.detection(frame)
        data = at_detect.get_point(tags)
        print(data)
        cv.imshow("image",at_detect.darw(frame,tags))
        key = cv.waitKey(10)
        if key == 27:
            break
    cv.destroyAllWindows()
    cap.release()


